import argparse
import warnings

import joblib
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler,LabelEncoder,OneHotEncoder
from sklearn.metrics import roc_auc_score,classification_report
from imblearn.over_sampling import SMOTE
import optuna
from optuna.samplers import TPESampler
from lightgbm import LGBMClassifier
from lightgbm.callback import early_stopping
import matplotlib.pyplot as plt
from xgboost import XGBClassifier
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


class PowerLoadModel:
    def __init__(self,path):
        # 获取数据源
        self.data_source = pd.read_csv(path)

def ana_data(data):
    """
    用来分析数据
    :param data:
    :return:
    """
    ana_data = data.copy()
    # print(ana_data.value_counts())  # 0 代表没有没有离职 1 代表离职
    # ana_data.info()
    # 加载数据（替换为你的数据路径）

    # 识别分类变量（根据业务理解，以下为典型分类特征）
    cat_cols = [
        "BusinessTravel", "Department", "EducationField", "Gender",
        "JobRole", "MaritalStatus", "Over18", "OverTime"
    ]

    # 对分类变量进行标签编码（将文字类别转为数值）
    le = LabelEncoder()
    for col in cat_cols:
        ana_data[col] = le.fit_transform(ana_data[col])

    # 对“离职情况（Attrition）”也进行编码
    ana_data["Attrition"] = le.fit_transform(ana_data["Attrition"])
    # 计算相关性矩阵，并提取与Attrition的相关性
    corr =ana_data.corr()["Attrition"].sort_values(ascending=False).to_frame()
    plt.figure(figsize=(8, 10))  # 设置图的大小
    sns.heatmap(
        corr,
        annot=True,  # 显示相关系数数值
        cmap="coolwarm",  # 颜色映射（红-正相关，蓝-负相关）
        vmin=-1, vmax=1  # 相关性范围（-1到1）
    )
    plt.title("特征与离职情况的相关性热力图")
    plt.tight_layout()  # 自动调整布局，避免标签重叠
    plt.show()
def feature_engineering(data):
    """
    给定数据源，进行特征处理 并返回
    :param data:
    :return:
    """
    feature_data = data.copy()
    # feature_data.info()
    # 根据上面的热力图来选取特征 先对一些object 进行one-hot处理
    one_data = feature_data[["Attrition","OverTime","MaritalStatus","DistanceFromHome","JobRole","Department","PerformanceRating","Gender",
                "JobInvolvement",'JobSatisfaction','StockOptionLevel','YearsWithCurrManager',
                'YearsInCurrentRole','MonthlyIncome','YearsAtCompany','JobLevel','Age','TotalWorkingYears']]
    feature_data = pd.get_dummies(one_data)
    feature_columns = feature_data.columns
    # print(feature_data.info())
    # print(feature_columns)
    return feature_data,feature_columns

def model_train(data,feature):
    """
    训练模型
    :param data: 数据
    :param feature: 特征
    :return:
    """
    x = data.iloc[:,1:]
    y = data["Attrition"]

    x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.2,random_state=1113)
    model =XGBClassifier(
        n_estimators=164,
        learning_rate= 0.030031526805699912,
        max_depth=7
        )
    model.fit(x_train,y_train)
    y_pre=model.predict(x_test)
    y_pred = model.predict_proba(x_test)[:, 1]
    joblib.dump(model, '../model/model.pkl')
    print('模型保存成功')
    print(f"训练集AUC:{round(roc_auc_score(y_test,y_pred)*100,2)}%")
    print(classification_report(y_test,y_pre))



if __name__ == '__main__':
    pm = PowerLoadModel('../data/train.csv')
    # print(pm.data_source)
    # ana_data(pm.data_source)
    feature_data,feature_columns=feature_engineering(pm.data_source)
    model_train(feature_data,feature_columns)
    # print(feature_columns)